import torch
device_gpu = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
device = torch.device("cpu")
# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
class ObservationBuffer:
    def __init__(self, num_envs, num_obs):
        self.num_envs = num_envs
        self.num_obs = num_obs

        # 用于存放obs，具体的每个细节数据存放后续再说
        self.obs_buf = torch.zeros(self.num_envs, self.num_obs, dtype=torch.float,device=device)
        self.depth_latent_and_yaw = torch.zeros(self.num_envs, 34, dtype=torch.float, device=device_gpu)
        self.obs_buf[:, 11] = 1
        self.obs_buf[:, 10:11] = 1
        self.default_pos = torch.tensor([ 0.0, 0.9, -1.8,
                        0.0, 0.9, -1.8, 
                        0.0, 1.12, -1.8,
                        0.0, 1.12, -1.8, ], dtype=torch.float32, device=device)
        """
        常量定义
        """
        self.dof_pos = 1.0
        self.dof_vel = 0.05
        self.action_scale = 0.25
        self.control_type = "P"
        self.stiffness = 40.0  # [N*m/rad]
        self.damping = 1.0  # [N*m*s/rad]

        self.torque_limit = 12.

    def insert_proprio(self, omega, acc, yaw, q, qd, contact):
        # 插入53个proprio数据
        # 上个action要放入进来
        self.obs_buf[:, : 3] = omega
        self.obs_buf[:, 3: 6] = acc
        self.obs_buf[:, 6: 8] = yaw
        # 关节角度和角速度
        self.obs_buf[:, 13: 25] = (q - self.default_pos) * self.dof_pos
        self.obs_buf[:, 25: 37] = qd * self.dof_vel
        # 添加触地数据
        self.obs_buf[:, 49: 53] = contact

    # 参数含义：上一个actions
    def insert(self, actions, priprio_last):
        # history要先向左移动然后放进来
        self.obs_buf[:, 37: 49] = actions
        # 这个是移动历史数据
        self.obs_buf[:, self.num_obs-530: self.num_obs - 53] = self.obs_buf[:,self.num_obs-530 + 53: self.num_obs].clone()
        self.obs_buf[:, self.num_obs - 53:] = priprio_last


    """
    FL -> FR -> RL -> RR
    """


    def compute_torques(self, actions, dof_pos, dof_vel):
        """Compute torques from actions.
            Actions can be interpreted as position or velocity targets given to a PD controller, or directly as scaled torques.
            [NOTE]: torques must have the same dimension as the number of DOFs, even if some DOFs are not actuated.

        Args:
            actions (torch.Tensor): Actions

        Returns:
            [torch.Tensor]: Torques sent to the simulation
        """
        # pd controller
        actions_scaled = actions * self.action_scale
        torques = (
            self.stiffness * (actions_scaled + self.default_pos - dof_pos)
            - self.damping * dof_vel
        )

        return torch.clip(torques, -self.torque_limit, self.torque_limit)